{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Relay PrintIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[21:04:47] /workspace/tvm/src/ir/transform.cc:655: PrintIR():\n",
      "#[version = \"0.0.5\"]\n",
      "def @main(%x: Tensor[(1, 2, 3), float32] /* ty=Tensor[(1, 2, 3), float32] */) -> Tensor[(1, 2, 3), float32] {\n",
      "  %0 = add(%x, %x) /* ty=Tensor[(1, 2, 3), float32] */;\n",
      "  multiply(%0, 2f /* ty=float32 */) /* ty=Tensor[(1, 2, 3), float32] */\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import tvm\n",
    "from tvm import relay\n",
    "\n",
    "shape = (1, 2, 3)\n",
    "tp = relay.TensorType(shape, \"float32\")\n",
    "x = relay.var(\"x\", tp)\n",
    "y = relay.add(x, x)\n",
    "y = relay.multiply(y, relay.const(2, \"float32\"))\n",
    "func = relay.Function([x], y)\n",
    "\n",
    "seq = tvm.transform.Sequential(\n",
    "    [\n",
    "        relay.transform.InferType(),\n",
    "        relay.transform.FoldConstant(),\n",
    "        tvm.transform.PrintIR(),\n",
    "        relay.transform.DeadCodeElimination(),\n",
    "    ]\n",
    ")\n",
    "\n",
    "mod = tvm.IRModule({\"main\": func})\n",
    "with tvm.transform.PassContext(opt_level=3):\n",
    "    mod = seq(mod)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = capfd.readouterr().err\n",
    "\n",
    "assert \"PrintIR\" in out\n",
    "assert \"multiply\" in out"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.13 ('py38': conda)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "28558e8daad512806f5c536a1a04c119185f99f65b79002708a12162d02a79c7"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
